package us.luosl.hadoop.mr.io;

import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.IOUtils;
import org.apache.hadoop.io.IntWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;

import java.io.IOException;
import java.io.OutputStream;

/**
 * Created by luosl on 2018/9/3.
 */
public class CustomOutputFormat extends FileOutputFormat{
    @Override
    public RecordWriter getRecordWriter(TaskAttemptContext job) throws IOException, InterruptedException {
        Path outPath = getOutputPath(job);
        FileSystem fs = outPath.getFileSystem(job.getConfiguration());
        Path nameCountPath = new Path(outPath, "name.txt");
        Path ageCountPath = new Path(outPath, "age.txt");
        String separator = job.getConfiguration().get("CustomRecordWriter.separator", "\001");
        return new CustomRecordWriter(fs.create(nameCountPath), fs.create(ageCountPath), separator);
    }

    public static class CustomRecordWriter extends RecordWriter<Text, IntWritable>{
        private String separator;
        private OutputStream nameOut;
        private OutputStream ageOut;

        public CustomRecordWriter(OutputStream nameOut, OutputStream ageOut, String separator) {
            this.separator = separator;
            this.nameOut = nameOut;
            this.ageOut = ageOut;
        }

        @Override
        public void write(Text key, IntWritable value) throws IOException, InterruptedException {
            String[] arr = key.toString().split("\001");
            if(arr.length > 1){
                if("name".equals(arr[0])){
                    write(nameOut, arr[1], value.get());
                }else if("age".equals(arr[0])){
                    write(ageOut, arr[1], value.get());
                }
            }
        }

        private void write(OutputStream out, String key, int value)throws IOException, InterruptedException{
            out.write(key.getBytes("UTF-8"));
            out.write(separator.getBytes("UTF-8"));
            out.write(Integer.toString(value).getBytes());
            out.write("\n".getBytes());
        }

        @Override
        public void close(TaskAttemptContext context) throws IOException, InterruptedException {
            IOUtils.closeStream(nameOut);
            IOUtils.closeStream(ageOut);
        }
    }
}
